# 1、导包、配置绘图字体
import os
import pandas as pd
import matplotlib.pyplot as plt
import datetime

from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler

from utils.log import Logger
from utils.common import data_preprocessing
from xgboost import XGBClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report, roc_auc_score
from sklearn.model_selection import GridSearchCV

import joblib
import seaborn as sns

# 解决图表中文乱码问题
plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['font.size'] = 15
# 忽略警告
import warnings

warnings.filterwarnings("ignore")
warnings.simplefilter(action='ignore', category=FutureWarning)


# 2、定义模型类，配置日志，获取数据源
class TalentLossModel(object):
    def __init__(self, filename):
        logfile_name = "train_" + datetime.datetime.now().strftime('%Y%m%d%H%M%S')
        self.logfile = Logger('../', logfile_name).get_logger()
        self.data_source = data_preprocessing(filename)


# 3、数据分析
def ana_data(data):
    data = data.copy(deep=True)
    # 1.数据整体情况
    data.info()
    print(data.head())
    print(data.columns)
    plt.pie(data.Attrition.value_counts(), labels=['在职', '离职'], autopct='%.2f%%')
    plt.title('员工离职比例')
    plt.show()
    # plt.savefig('人才离职情况分析图.png')

    # 数值型数据的可视化分析
    # 与家的距离与离职的相关性
    sns.histplot(x='DistanceFromHome', hue='Attrition', data=data, kde=True)
    plt.xlabel('与家的距离')  # 设置x轴名称
    plt.ylabel('人数')  # 设置y轴名称
    plt.title('与家的不同距离的离职情况')
    # plt.show()
    # 绩效评级与离职的相关性
    sns.countplot(x='PerformanceRating', hue='Attrition', data=data)
    plt.xlabel('绩效评级')  # 设置x轴名称
    plt.ylabel('人数')  # 设置y轴名称
    plt.title('不同绩效评级下的离职情况')
    # plt.show()
    # 工资涨幅百分比与离职的相关性
    sns.countplot(x='PercentSalaryHike', hue='Attrition', data=data)
    plt.xlabel('工资涨幅百分比')  # 设置x轴名称
    plt.ylabel('人数')  # 设置y轴名称
    plt.title('不同工资涨幅百分比的离职情况')
    # plt.show()
    # # 将Gender、Over18、OverTime列的数据转换为数值特征
    data['Gender'] = data['Gender'].map({'Male': 1, 'Female': 0})
    data['Over18'] = data['Over18'].map({'Y': 1, 'N': 0})
    data['OverTime'] = data['OverTime'].map({'Yes': 1, 'No': 0})

    # 筛选出数值型列的数据
    numerical_data = data.select_dtypes(include=['int64', 'float64'])
    # print(numerical_data)   # [1100 rows x 23 columns]
    # 绘制数值型列与离职的相关性
    plt.figure(figsize=(30, 25))
    plot = 1
    for var in numerical_data:
        plt.subplot(9, 3, plot)
        sns.distplot(data[var])
        plot += 1
        plt.tight_layout()
    # plt.show()

    # 非数值型数据的可视化分析
    # 加班情况与离职的相关性
    sns.countplot(x='OverTime', hue='Attrition', data=data)
    plt.xlabel('加班情况')
    plt.ylabel('人数')
    plt.title('不同加班情况的离职情况')
    plt.tight_layout()
    plt.show()
    # 工作角色与离职的相关性
    plt.figure(figsize=(20, 15))
    sns.countplot(x='JobRole', hue='Attrition', data=data)
    plt.xticks(rotation=15)
    plt.xlabel('工作角色')
    plt.ylabel('人数')
    plt.title('不同工作角色的离职情况')
    plt.tight_layout()
    plt.show()
    # 部门与离职的相关性
    sns.countplot(x='Department', hue='Attrition', data=data)
    plt.xlabel('部门')
    plt.ylabel('人数')
    plt.title('不同部门的离职情况')
    plt.tight_layout()
    plt.show()
    # 结果：Sales 部门离职人数较多；Research & Development 部门离职人数较少
    # 教育领域与离职的相关性
    sns.countplot(x='EducationField', hue='Attrition', data=data)
    plt.xticks(rotation=25)
    plt.xlabel('教育领域')
    plt.ylabel('人数')
    plt.title('不同教育领域的离职情况')
    plt.tight_layout()
    plt.show()

    # 在原始数据中筛选出非数值型列的数据
    categorical_data = data.select_dtypes(include=['object'])
    # categorical_data.info()    # [1100 rows x 8 columns]
    # 绘制在原始数据中非数值型列与离职的相关性
    plt.figure(figsize=(30, 25))
    plot = 1
    for feature_name in categorical_data:
        plt.subplot(4, 2, plot)
        sns.countplot(x=feature_name, hue='Attrition', data=data)
        plot += 1
        plt.xticks(rotation=15)
        plt.tight_layout()
    plt.show()


# 4、特征工程（重点）
def feature_engineering(data, logger):
    logger.info('=======================开始进行特征工程处理==============================')
    result = data.copy(deep=True)

    logger.info('=======================开始提取特征==============================')
    # 1、提取出特征
    # 非数值型数据
    cat_data = result.select_dtypes(include=['object'])
    for column in cat_data.columns:
        result[column] = pd.factorize(result[column])[0]
    # print("当前数据集中各列的数据类型：\n", result.dtypes)
    full_data = result.copy(deep=True)
    num_data = result.select_dtypes(exclude=['object'])
    num_crr = num_data.corr(method='spearman')['Attrition']
    print(num_crr.sort_values(ascending=False))
    top_features = num_crr.dropna().sort_values(ascending=False).index[1:]
    """
    .corr(method='spearman')：计算所有列之间的 Spearman 等级相关系数，衡量它们与“Attrition”（离职）之间的单调关系
    [1:]：切片操作，去掉第一个元素（即 'Attrition' 自身）
    .dropna()：删除结果为NaN的数据
    """
    feature_data = result[top_features]
    # data = data.drop(columns=['Over18', 'StandardHours'])
    return feature_data, full_data


# 5、模型训练、评价与保存
def model_train(data, feature, logger):
    logger.info('=======================开始进行模型训练==============================')
    x_data = data[feature]
    y_data = data['Attrition']
    x_train, x_test, y_train, y_test = train_test_split(x_data, y_data, test_size=0.2, random_state=1)
    # 特征处理
    ss = StandardScaler()
    x_train = ss.fit_transform(x_train)
    x_test = ss.transform(x_test)
    # 模型训练
    xgb = XGBClassifier(n_estimators=150, max_depth=6, learning_rate=0.1)
    xgb.fit(x_train, y_train)
    y_pred_test = xgb.predict(x_test)
    print(f'xgb模型的准确率为：{accuracy_score(y_test, y_pred_test)}')
    print('xgb模型的分类报告：\n', classification_report(y_test, y_pred_test))
    print('xgb模型的AUC：', roc_auc_score(y_test, y_pred_test))

    # 使用逻辑回归模型
    model = LogisticRegression(C=0.1, class_weight='balanced', penalty='l1', solver='saga')
    model.fit(x_train, y_train)
    y_pred = model.predict(x_test)
    print('逻辑回归模型的准确率：', accuracy_score(y_test, y_pred))
    print('逻辑回归模型的分类报告：\n', classification_report(y_test, y_pred))
    print('逻辑回归模型的AUC:', roc_auc_score(y_test, y_pred))
    logger.info('=======================模型训练完成============================')
    joblib.dump(xgb, '../model/ljhg.pkl')


# 6、模型训练主函数
if __name__ == '__main__':
    # 加载数据
    input_file = os.path.join('../data', 'train.csv')

    model = TalentLossModel(input_file)

    # 数据可视化分析
    # ana_data(model.data_source)

    # 特征工程
    processed_data, full_data = feature_engineering(model.data_source, model.logfile)
    print("processed_data是特征列：\n", processed_data.columns)
    print("feature_cols是完整数据列：\n", full_data.columns)
    # processed_data是处特征列，feature_cols是完整数据列
    # 模型训练、评价与保存
    model_train(full_data, processed_data.columns, model.logfile)
